# -*- encoding: utf-8 -*-
import ID3
import Gini
import Error
from collections import Counter


def create_training_data():
    """创建训练数据"""
    dataset = [
        ['Sunny', 'Hot', 'High', 'Weak', 'No'],
        ['Sunny', 'Hot', 'High', 'Strong', 'No'],
        ['Overcast', 'Hot', 'High', 'Weak', 'Yes'],
        ['Rain', 'Mild', 'High', 'Weak', 'Yes'],
        ['Rain', 'Cool', 'Normal', 'Weak', 'Yes'],
        ['Rain', 'Cool', 'Normal', 'Strong', 'No'],
        ['Overcast', 'Cool', 'Normal', 'Strong', 'Yes'],
        ['Sunny', 'Mild', 'High', 'Weak', 'No'],
        ['Sunny', 'Cool', 'Normal', 'Weak', 'Yes'],
        ['Rain', 'Mild', 'Normal', 'Weak', 'Yes'],
        ['Sunny', 'Mild', 'Normal', 'Strong', 'Yes'],
        ['Overcast', 'Mild', 'High', 'Strong', 'Yes'],
        ['Overcast', 'Hot', 'Normal', 'Weak', 'Yes'],
        ['Rain', 'Mild', 'High', 'Strong', 'No'],
    ]
    labels = ['Outlook', 'Temperature', 'Humidity', 'Wind']
    return dataset, labels


def create_id3_decision_tree(dataset, labels):
    """构建id3决策树"""
    # 递归终止条件1：传入数据集的分类只有一种
    if len(list(set([data[-1] for data in dataset]))) == 1:
        return dataset[0][-1]
    # 递归终止条件2：全部特征已经用完，则选择数量最多的分类最为结果返回
    if len(dataset[0]) == 1:
        return Counter([data[-1] for data in dataset]).most_common(1)[0][0]

    best_feature = ID3.find_best_feature(dataset)
    best_feature_label = labels[best_feature]
    my_tree = {best_feature_label: {}}
    best_feature_set = set([data[best_feature] for data in dataset])
    sub_labels = list(labels)
    del sub_labels[best_feature]
    for value in best_feature_set:
        my_tree[best_feature_label][value] = create_id3_decision_tree(ID3.split_dataset(dataset, best_feature, value),
                                                                      sub_labels)
    return my_tree


def create_gini_decision_tree(dataset, labels):
    """构建gini split决策树"""
    # 递归终止条件1：传入数据集的分类只有一种
    if len(list(set([data[-1] for data in dataset]))) == 1:
        return dataset[0][-1]
    # 递归终止条件2：全部特征已经用完，则选择数量最多的分类最为结果返回
    if len(dataset[0]) == 1:
        return Counter([data[-1] for data in dataset]).most_common(1)[0][0]

    best_feature = Gini.find_best_feature(dataset)
    best_feature_label = labels[best_feature]
    my_tree = {best_feature_label: {}}
    best_feature_set = set([data[best_feature] for data in dataset])
    sub_labels = list(labels)
    del sub_labels[best_feature]
    for value in best_feature_set:
        my_tree[best_feature_label][value] = create_gini_decision_tree(Gini.split_dataset(dataset, best_feature, value),
                                                                      sub_labels)
    return my_tree


def create_error_decision_tree(dataset, labels):
    """构建misclassification Error决策树"""
    # 递归终止条件1：传入数据集的分类只有一种
    if len(list(set([data[-1] for data in dataset]))) == 1:
        return dataset[0][-1]
    # 递归终止条件2：全部特征已经用完，则选择数量最多的分类最为结果返回
    if len(dataset[0]) == 1:
        return Counter([data[-1] for data in dataset]).most_common(1)[0][0]

    best_feature = Error.find_best_feature(dataset)
    best_feature_label = labels[best_feature]
    my_tree = {best_feature_label: {}}
    best_feature_set = set([data[best_feature] for data in dataset])
    sub_labels = list(labels)
    del sub_labels[best_feature]
    for value in best_feature_set:
        my_tree[best_feature_label][value] = create_gini_decision_tree(Error.split_dataset(dataset, best_feature, value),
                                                                      sub_labels)
    return my_tree


def classify(input_tree, test_labels, test_value):
    """测试决策树分类"""
    first_feature = list(input_tree.keys())[0]
    second_dict = input_tree[first_feature]
    first_index = test_labels.index(first_feature)
    first_value = test_value[first_index]
    first_feature_value = second_dict[first_value]
    if first_feature_value in ('Yes', 'No'):
        return first_feature_value
    result = classify(first_feature_value, test_labels, test_value)
    return result


if __name__ == '__main__':
    dataset, labels = create_training_data()
    test_labels, test_value = ['Outlook', 'Temperature', 'Humidity', 'Wind'], ['Overcast', 'Cool', 'Normal', 'Strong']

    my_tree = create_id3_decision_tree(dataset, labels)
    print('ID3决策树：', my_tree)
    result = classify(my_tree, test_labels, test_value)
    print('ID3算法结果', result)
    my_tree = create_gini_decision_tree(dataset, labels)
    print('Gini split决策树：', my_tree)
    result = classify(my_tree, test_labels, test_value)
    print('Gini split算法结果', result)
    my_tree = create_error_decision_tree(dataset, labels)
    print('Misclassification Error决策树：', my_tree)
    result = classify(my_tree, test_labels, test_value)
    print('Misclassification Error算法结果', result)
